/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"
#include <iostream>

using namespace std;

size_t s_pid = sizeof(PageId);
size_t s_rid = sizeof(RecordId);
size_t s_key = sizeof(int);
int s_entry = s_rid + s_key;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
    RC rc;

	// Try to open PageFile and return error if open caused error
	if ((rc = pf.open(indexname, mode)) < 0) {
		return rc;
	}
	
	// Create buffer to read from PageFile
	char buf[PageFile::PAGE_SIZE];
	memset(&buf, 0, PageFile::PAGE_SIZE);
	rc = pf.read(0, buf);
	
	// Get info for rootPid and treeHeight from page file
	memcpy(&rootPid, &buf[0], s_pid);
	memcpy(&treeHeight, &buf[s_pid], s_key);
	
	// Write info back to PageFile
	pf.write(0, buf);
	
	return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
	RC rc;
	
	// Write rootPid and treeHeight to PageFile
	char buf[PageFile::PAGE_SIZE];
	memset(buf, 0, PageFile::PAGE_SIZE);
	memcpy(&buf[0], &rootPid, s_pid);
	memcpy(&buf[s_pid], &treeHeight, s_key);
	
	if ((rc = pf.write(0, buf)) < 0) {
		return rc;
	}
	
	// Close the PageFile
	if ((rc = pf.close()) < 0) {
		return rc;
	}
	
	return 0;
}

int BTreeIndex::insertHelper(int key, const RecordId& rid, int height, PageId& pid)
{
	int rc;
	int sibling_key = 0;
	
	// Inserting into an empty tree
	if (treeHeight == 0) {
		// Create a new root node
		BTNonLeafNode root;
		PageId root_pid = pf.endPid();
		
		if ((rc = root.read(root_pid, pf)) < 0 && rc != RC_INVALID_PID)
			return rc;
			
		// Create two new leafs
		BTLeafNode n1, n2;
		PageId n1_pid = pf.endPid() + 1;
		if ((rc = n1.read(n1_pid, pf)) < 0 && rc != RC_INVALID_PID)
			return rc;
			
		PageId n2_pid = pf.endPid() + 2;
		if ((rc = n2.read(n2_pid, pf)) < 0 && rc != RC_INVALID_PID)
			return rc;
			
		// Connect leafs & insert key
		n1.setNextNodePtr(n2_pid);
		n2.insert(key, rid);
		
		// Insert into root
		root.initializeRoot(n1_pid, key, n2_pid);
		// Write into nodes
		root.write(root_pid, pf);
		n1.write(n1_pid, pf);
		n2.write(n2_pid, pf);
		
		// Set rootPid, increase treeheight and return success
		rootPid = root_pid;
		treeHeight++;
		return 0;
	}
	// Insert into a leaf node
	else if (height == treeHeight) {
		// Create a new leaf node
		BTLeafNode leaf_node;
		
		if ((rc = leaf_node.read(pid, pf)) < 0) {
			return rc;
		}
		
		// Insert into leaf node
		// If node is full, then make a new sibling node, insert, and split
		if ((rc = leaf_node.insert(key, rid)) == RC_NODE_FULL) {
			// Make sibling node
			BTLeafNode sibling_node;
			PageId new_pid = pf.endPid();
			
			if ((rc = sibling_node.read(new_pid, pf)) < 0 && rc != RC_INVALID_PID) {
				return rc;
			}
			
			// Insert key and split
			leaf_node.insertAndSplit(key, rid, sibling_node, sibling_key);
			// Set next node info
			PageId next = leaf_node.getNextNodePtr();
			sibling_node.setNextNodePtr(next);
			leaf_node.setNextNodePtr(new_pid);
			
			leaf_node.write(pid, pf);
			sibling_node.write(new_pid, pf);
			
			// Set return pid for recursive call and sibling_key for child_ptr
			pid = new_pid;
			return sibling_key;
		}
		// If node isn't full, return success
		else {
			leaf_node.write(pid, pf);
			return 0;
		}
	}
	// Insert into a nonleaf node
	else {
		// Create new nonleaf node
		BTNonLeafNode nonleaf_node;
		PageId child_pid;
		
		if ((rc = nonleaf_node.read(pid, pf)) < 0) {
			return rc;
		}
		nonleaf_node.locateChildPtr(key, child_pid);
		
		// Do a recursive insert into tree with an increase height
		// If tree is full
		if ((sibling_key = insertHelper(key, rid, height+1, child_pid)) > 0) {
			// Insert into nonleaf node
			// If nonleaf node is full
			if ((rc = nonleaf_node.insert(sibling_key, child_pid)) == RC_NODE_FULL) {
				// Create a sibling node and split
				BTNonLeafNode sibling_node;
				PageId new_pid = pf.endPid();
			
				if ((rc = sibling_node.read(new_pid, pf)) < 0 && rc != RC_INVALID_PID) {
					return rc;
				}
			
				// Insert key and split
				nonleaf_node.insertAndSplit(key, child_pid, sibling_node, sibling_key);
				nonleaf_node.write(pid, pf);
				sibling_node.write(new_pid, pf);
			
				// Current node is root
				if (pid != rootPid) {
					// Set return pid for recursive call and sibling_key for child_ptr
					pid = new_pid;
					return sibling_key;
				}
				// Not root node
				else {
					// Create a new root node
					BTNonLeafNode root_node;
					PageId old_pid = rootPid;
					// Set new root pid
					rootPid = pf.endPid();
					
					if ((rc = root_node.read(rootPid, pf)) < 0 && rc != RC_INVALID_PID) {
						return rc;
					}
					
					// Init new root
					root_node.initializeRoot(old_pid, sibling_key, new_pid);
					root_node.write(rootPid, pf);
					
					// Increment tree height and return
					treeHeight++;
					return 0;
				}
			}
			// Nonleaf node isn't full
			else {
				nonleaf_node.write(pid, pf);
				return 0;
			}
		}
		// If the tree isn't full
		else {
			return 0;
		}
	}
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	RC rc;
    if ((rc = insertHelper(key, rid, 0, rootPid)) < 0 )
		return rc;
		
	return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	// Set rootPid to next PageId
	PageId next_pid = rootPid;
	BTNonLeafNode node;
	
	// Search through non-leaf nodes of tree starting from bottom up
	for (int i = treeHeight; i > 0; i--) {
		node.read(next_pid, pf);
		node.locateChildPtr(searchKey, next_pid);
	}
	
	// Search leaf nodes
	BTLeafNode leaf;
	int eid;
	
	leaf.read(next_pid, pf);
	leaf.locate(searchKey, eid);
	
	// Set cursor fields
	cursor.pid = next_pid;
	cursor.eid = eid;
	
	// Return 0 for success
    return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	RC rc;
	
	// Create a buffer and read from cursor point into buffer
	char buf[PageFile::PAGE_SIZE];
	memset(buf, 0, PageFile::PAGE_SIZE);
	
	if ((rc = pf.read(cursor.pid, buf)) < 0) {
		return rc;
	}
	
	// Get index and read info from buffer
	int index = cursor.eid * s_entry;
	
	memcpy(&rid, &buf[index], s_rid);
	key = *((int*) &buf[index + s_rid]);
	
	// Adjust cursor eid
	if (cursor.eid < PageFile::PAGE_SIZE/s_entry - 1 && key > 0) {
		cursor.eid += 1;
	}
	else {
		memcpy(&cursor.pid, &buf[PageFile::PAGE_SIZE - s_pid], s_pid);
		if (cursor.pid <= 0)
			return RC_INVALID_PID;
		cursor.eid = 0;
	}
	
	// Return 0 for success
    return 0;
}
